import math
import struct





# 姿态角全局变量
pitch = 0.0
roll = 0.0
yaw = 0.0

# 偏差和协方差矩阵
pitch_bias = 0.0
roll_bias = 0.0

P_pitch = [[0.0, 0.0], [0.0, 0.0]]
P_roll  = [[0.0, 0.0], [0.0, 0.0]]

# 卡尔曼滤波参数
Q_angle = 0.003
Q_bias = 0.003
R_measure = 0.003


#快速平方根倒数
def Q_rsqrt(number):
    """
    快速计算 1/√number 的近似值。

    参数:
        number: 正浮点数

    返回:
        1/√number 的近似值
    """
    threehalfs = 1.5

    x2 = number * 0.5
    y  = number
    # 将浮点数的二进制表示转换为整数
    packed_y = struct.pack('f', y)
    i = struct.unpack('I', packed_y)[0]
    # 魔法常数，见于 Quake III 的实现
    i = 0x5f3759df - (i >> 1)
    # 将整数转换回浮点数
    packed_i = struct.pack('I', i)
    y = struct.unpack('f', packed_i)[0]
    # 牛顿迭代，1次迭代通常就已经足够了
    y = y * (threehalfs - (x2 * y * y))
    return y

def turn_imu_data(imu_data):
    """
    将陀螺仪数据转换成真实物理数据
    
    参数：
        imu_data: 陀螺仪get()数据
        
    返回：
        真实数据构成的元组
    """
    ax = imu_data[0]*0.244/1000
    ay = imu_data[1]*0.244/1000
    az = imu_data[2]*0.244/1000
    gx = imu_data[3]*0.07*0.0174
    gy = imu_data[4]*0.07*0.0174
    gz = imu_data[5]*0.07*0.0174
    
    return (ax,ay,az,gx,gy,gz)

def kalman_filter(imu_data_turn, dt):
    """
    输入：六轴传感器 + 时间间隔
    输出：无（直接更新全局变量 pitch、roll、yaw）
    """
    global pitch, roll, yaw
    global pitch_bias, roll_bias
    global P_pitch, P_roll
    ax, ay, az, gx, gy, gz = imu_data_turn
    # ---------- 加速度估计角度 ----------
    inv_sqrt_yz = Q_rsqrt(ay * ay + az * az)
    acc_roll = math.degrees(math.atan(-ax * inv_sqrt_yz))
    acc_pitch = math.degrees(math.atan2(ay, az))

    # ---------- Pitch 轴更新 ----------
    rate = gx - pitch_bias
    pitch += dt * rate

    P = P_pitch
    P[0][0] += dt * (dt * P[1][1] - P[0][1] - P[1][0] + Q_angle)
    P[0][1] -= dt * P[1][1]
    P[1][0] -= dt * P[1][1]
    P[1][1] += Q_bias * dt

    S = P[0][0] + R_measure
    K0 = P[0][0] / S
    K1 = P[1][0] / S
    y_pitch = acc_pitch - pitch
    pitch += K0 * y_pitch
    pitch_bias += K1 * y_pitch

    P00_temp = P[0][0]
    P01_temp = P[0][1]
    P[0][0] -= K0 * P00_temp
    P[0][1] -= K0 * P01_temp
    P[1][0] -= K1 * P00_temp
    P[1][1] -= K1 * P01_temp

    # ---------- Roll 轴更新 ----------
    rate = gy - roll_bias
    roll += dt * rate

    P = P_roll
    P[0][0] += dt * (dt * P[1][1] - P[0][1] - P[1][0] + Q_angle)
    P[0][1] -= dt * P[1][1]
    P[1][0] -= dt * P[1][1]
    P[1][1] += Q_bias * dt

    S = P[0][0] + R_measure
    K0 = P[0][0] / S
    K1 = P[1][0] / S
    y_roll = acc_roll - roll
    roll += K0 * y_roll
    roll_bias += K1 * y_roll

    P00_temp = P[0][0]
    P01_temp = P[0][1]
    P[0][0] -= K0 * P00_temp
    P[0][1] -= K0 * P01_temp
    P[1][0] -= K1 * P00_temp
    P[1][1] -= K1 * P01_temp

    # ---------- Yaw 直接积分 ----------
    yaw += gz * dt
    
    return pitch,roll,yaw,gz,acc_pitch

def get_kalman_imu(imu_data):
    imu_data_turn = turn_imu_data(imu_data)
    pitch,roll,yaw,gz,ac_p = kalman_filter(imu_data_turn,0.01)
    
    return pitch,roll,yaw,gz,ac_p